#include <string.h>
#include <unistd.h>
#include <sys/wait.h>
#include <thread>
#include <vector>
#include <future>
#include <string>
#include <functional>
#include <iostream>
 #include <utility> // std::forward

template <typename F, typename... Args>
int launch_multi_task(int world_size, const F&& func, Args&&... args) {
    std::vector<pid_t> pids(world_size);
    for (int rank = 0; rank < world_size; rank ++) {
        pid_t pid = fork();
        if (pid == 0) {
            std::string str_rank = std::to_string(rank);
            std::string str_world_size = std::to_string(world_size);
            setenv("MASTER_ADDR", "127.0.0.1", 1);
            setenv("MASTER_PORT", "9960", 1);
            setenv("RANK", str_rank.c_str(), 1);
            setenv("WORLD_SIZE", str_world_size.c_str(), 1);
            return func(std::forward<Args>(args)...);
        } else if (pid < 0) {
            std::cerr << "fork failed, reason: " << strerror(errno) << std::endl;
            return errno;
        }
        pids[rank] = pid;
    }
    int status = 0;
    for (int i = 0; i < world_size; i ++) {
        waitpid(pids[i], &status, 0);
        if (status != 0) {
            for (int j = i + 1; j < world_size; j ++) {
                kill(pids[j], SIGSTOP);
            }
            return status;
        }
    }
    return 0;
}

int add_function(int a, int b) {
    std::cout << "add_function result: " << (a + b) << std::endl;
    return 0;
}

int void_function() {
    std::cout << "I am a void funtion." << std::endl;
    return 0;
}

int main() {
    launch_multi_task(4, add_function, 2, 5);
    // launch_multi_task(3, void_function);
}
